#include "../Base.h"


THREAD_ENTRY threadOfHandleTable;
PROTECTEDENTRYINFO protectedEntryInfo;



//ȫ�־����
PVOID PidToHandleTableEntry(PHANDLE_TABLE p_handleTable,ULONG_PTR pid) {

	UINT32 type;
	ULONG_PTR index = 0;
	ULONG_PTR table_level_0 = 0,table_level_1 = 0,table_level_2 = 0;
	ULONG_PTR index_0 = 0, index_1 = 0, index_2 = 0;
	index = pid / 4;

	type = p_handleTable->TableCode & 0x3;
	
	switch (type)
	{

	case 0 :
	{

		table_level_0 = *(UINT32*)(p_handleTable->TableCode & ~0x3);
		index_0 = index % 512;
		break;
	}
	
	case 1:
	{
		index_0 = index / 512;
		index_1 = index_0 % 1024;
		table_level_1 = *(UINT32*)(p_handleTable->TableCode & ~0x3);
		table_level_0 = *(UINT32*)(table_level_1 + index_1 * 4);
		break;
	}
	case 2:
	{
		index_0 = index / 512;
		index_1 = index_0 / 1024;
		index_2 = index_1 % 1024;
		table_level_2 = *(UINT32*)(p_handleTable->TableCode & ~0x3);
		table_level_1 = *(UINT32*)(table_level_2 + index_2 * 4);
		table_level_0 = *(UINT32*)(table_level_1 + index_1 * 4);
		break;
	}
	default:
		break;
	}


	return (PVOID)(table_level_0 + index_0 * 8);

}


PVOID GetHandleTableEntryByProcessPid(ULONG_PTR pid) {
	PHANDLE_TABLE pHandleTable;
	ULONG_PTR table_index;
	pHandleTable  = (PHANDLE_TABLE)SearchPspCidTable();
	return PidToHandleTableEntry(pHandleTable,pid);
}



PVOID GetHandlTableEntryByProcessName(char* processName) {
	PEPROCESS process = FindProcessByNameInBaseName(processName);
	if (!process) {
		return NULL;
	}
	HANDLE pid = PsGetProcessId(process);
	return GetHandleTableEntryByProcessPid((ULONG_PTR)pid);
}

PVOID GetExpLookupHandleTableEntryAddress()
{

	UNICODE_STRING funcName;
	RtlInitUnicodeString(&funcName, L"ExEnumHandleTable");
	char buf[] = {
	0xE8, '*','*', '*','*',0x8B, 0xF0
	};
	PVOID funcAddress = MmGetSystemRoutineAddress(&funcName);
	PCH address = (PCH)SearchPatterFuzzy((PCH)funcAddress, 0x100, buf, sizeof(buf));
	address = (address + 5) + *(UINT32*)(address + 1);
	if (!address) return 0;
	return PVOID(address);
}

PVOID GetTableEntryByProcessId(UINT32 pid)
{
	pfnExpLookupHandleTableEntry MyExpLookupHandleTableEntry;
	PVOID pHandleTable = SearchPspCidTable();
	MyExpLookupHandleTableEntry = (pfnExpLookupHandleTableEntry)GetExpLookupHandleTableEntryAddress();
	if (!MyExpLookupHandleTableEntry) {
		return NULL;
	}
	
	return MyExpLookupHandleTableEntry((UINT32*) * (UINT32*)pHandleTable, pid);
}
PVOID GetTableEntryByProcessName(char* processName) {

	HANDLE pid;
	PEPROCESS process = FindProcessByNameInBaseName(processName);
	if (!process) {
		return NULL;
	}
	pid = PsGetProcessId(process);
	return GetTableEntryByProcessId((UINT32)pid);
}


//1.��ȡһ�����̵�˽�о����
//2.�����Ƿ������˽�о��
//3.����о��޸�Ȩ��

//˽�о����

PVOID GetPrivateTableEntryByProcessNameAndHandle(char* name,ULONG handle) {

	pfnExpLookupHandleTableEntry MyExpLookupHandleTableEntry;
	PEPROCESS_XP process = (PEPROCESS_XP)FindProcessByNameInBaseName(name);
	if (!process)
		return NULL;
	MyExpLookupHandleTableEntry = (pfnExpLookupHandleTableEntry)GetExpLookupHandleTableEntryAddress();
	if (!MyExpLookupHandleTableEntry) {
		return NULL;
	}
	return MyExpLookupHandleTableEntry((UINT32*)process->ObjectTable, handle);

}


BOOLEAN MyEnumHandleTableRoutine(
	IN PHANDLE_TABLE_ENTRY HandleTableEntry,
	IN HANDLE Handle,
	IN PVOID EnumParameter
) {

	PEPROCESS_XP process = (PEPROCESS_XP)EnumParameter;
	if ((PEPROCESS_XP)((HandleTableEntry->Value & ~7) +0x18) == process) {
		HandleTableEntry->GrantedAccess &=  ~(PROCESS_VM_READ | PROCESS_VM_WRITE);
	}
	return false;
}



VOID
threadRoutine(
	__in PVOID StartContext
) {


	while (threadOfHandleTable.isWorkThread) {

		PEPROCESS_XP process = (PEPROCESS_XP)StartContext;
		HANDLE pid = PsGetProcessId((PEPROCESS)process);
		for (int i = 0; i < 200000; i++) {
			if (pid == (HANDLE)i) continue;
			PEPROCESS_XP targetProcess;
			NTSTATUS status = PsLookupProcessByProcessId((HANDLE)i, (PEPROCESS*)&targetProcess);
			if (!NT_SUCCESS(status)) {
				continue;
			}
			HANDLE handle;
			ExEnumHandleTable(targetProcess->ObjectTable, MyEnumHandleTableRoutine, process, &handle);
			ObfDereferenceObject(targetProcess);
		}

		LARGE_INTEGER time;
		time.QuadPart = -30 * 1000 * 1000;
		KeDelayExecutionThread(KernelMode,FALSE,&time);
	}

}


PVOID DownRightHandleTableOfAllProcess(char* name)
{
	PEPROCESS_XP process = (PEPROCESS_XP)FindProcessByNameInBaseName(name);
	if (!process) {
		return NULL;
	}

	PsCreateSystemThread(&threadOfHandleTable.handle, THREAD_ALL_ACCESS, NULL, NULL, NULL, threadRoutine, process);
}

PVOID SaveProcessObject(PVOID object_header_header,OUT PVOID virtualAddress) {
#ifdef _WIN64
	ULONG size = 0x30;
#else
	ULONG size = 0x18;
#endif
	memcpy(virtualAddress, object_header_header, size + sizeof(EPROCESS_XP));
	return (PVOID)MmGetPhysicalAddress(virtualAddress).QuadPart;
}

BOOLEAN ChangeEprocessInfo(PEPROCESS_XP process,PVOID virtualAddress) {

	ULONG copySize = 0x20;
	process->UniqueProcessId = (void*)0x22222;
	char buf[16] = "test.name";
	memcpy(process->ImageFileName, buf, sizeof(buf));
	ULONG cr3 = process->Pcb.DirectoryTableBase[0];
	if (cr3 & 0x1000 == 0) {
		copySize = 0x1000;
	}
	LARGE_INTEGER cr3Address;
	cr3Address.LowPart = cr3 & 0x1000;
	PVOID cr3VirtuallAddress = MmMapIoSpace(cr3Address, copySize, MmNonCached);
	if (!cr3VirtuallAddress) return FALSE;
	memcpy(virtualAddress, cr3VirtuallAddress, copySize);
	process->Pcb.DirectoryTableBase[0] = MmGetPhysicalAddress(virtualAddress).QuadPart;
	return TRUE;
}

BOOLEAN ProtectPrivateHandleTableEntry(PVOID table_entry_address, ULONG handle)
{
	PVOID object_header_address,virtualAddress,physicalAddress,cr3_va,cr3_pa;
	if (!table_entry_address) return FALSE;
	object_header_address = (PVOID)(*(ULONG*)table_entry_address & ~0x7);
	virtualAddress = ExAllocatePoolWithTag(NonPagedPool, 0x1000, 'yang');
	physicalAddress = SaveProcessObject(object_header_address, virtualAddress);
	if (!physicalAddress) return FALSE;
	PEPROCESS_XP process = (PEPROCESS_XP)( & (((POBJECT_HEADER)virtualAddress)->Body));
	cr3_va = ExAllocatePoolWithTag(NonPagedPool, 0x1000, 'yang');
	
	if (ChangeEprocessInfo(process, cr3_va)) return FALSE;

	*(ULONG64*)table_entry_address = ((ULONG64)virtualAddress & ~0xFFFFFFFF00000007) | ((*(ULONG64*)table_entry_address) & 0xFFFFFFFF00000007);
	
}

BOOLEAN MyEnumHandleTableUpRightRoutine(
	IN PHANDLE_TABLE_ENTRY HandleTableEntry,
	IN HANDLE Handle,
	IN PVOID EnumParameter
) {

	PEPROCESS_XP process = (PEPROCESS_XP)EnumParameter;
	if ((PEPROCESS_XP)((HandleTableEntry->Value & ~7) + 0x18) == process) {
	
		ProtectPrivateHandleTableEntry(HandleTableEntry, (ULONG)Handle);
		HandleTableEntry->GrantedAccess |= PROCESS_ALL_ACCESS;
	}
	return false;
}

VOID
protectedEntryThreadRoutine(
	__in PVOID StartContext
) {

	PPROTECTEDENTRYINFO entryInfo = (PPROTECTEDENTRYINFO)StartContext;
	while (1) {
		HANDLE handle;
		ExEnumHandleTable(((PEPROCESS_XP)(entryInfo->protectedProcess))->ObjectTable, MyEnumHandleTableUpRightRoutine, entryInfo->targetProcess, &handle);

		LARGE_INTEGER time;
		time.QuadPart = -10 * 1000 * 1000;
		KeDelayExecutionThread(KernelMode, FALSE, &time);
	}
}

BOOLEAN ProtectPrivateHandleTableEntryByTargetProcessName(char* protectedProcessName, char* targetProcessName) {


	protectedEntryInfo.protectedProcess = FindProcessByNameInBaseName(protectedProcessName);
	if (!protectedEntryInfo.protectedProcess) return FALSE;
	protectedEntryInfo.targetProcess = FindProcessByNameInBaseName(targetProcessName);
	if (!protectedEntryInfo.targetProcess) return FALSE;
	PsCreateSystemThread(&threadOfHandleTable.handle, THREAD_ALL_ACCESS, NULL, NULL, NULL, protectedEntryThreadRoutine, &protectedEntryInfo);
	
}